#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/twist_stamped.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>

#include <tf2_ros/transform_listener.h>
#include <tf2_ros/buffer.h>
#include <eigen3/Eigen/Dense>
#include <cmath>
#include <memory>
#include <thread>

// for keyboard
#include <termios.h>
#include <unistd.h>
#include <fcntl.h>

#include <moveit/planning_scene_monitor/planning_scene_monitor.h>
#include <moveit_servo/servo_parameters.h>
#include <moveit_servo/servo.h>
#include <moveit/robot_state/robot_state.h>

using namespace std::chrono_literals;
static const rclcpp::Logger LOGGER = rclcpp::get_logger("force_servo_control_node");

// 获取键盘输入（非阻塞）
int getch()
{
  struct termios oldt, newt;
  int ch;
  int oldf;

  tcgetattr(STDIN_FILENO, &oldt);
  newt = oldt;
  newt.c_lflag &= ~(ICANON | ECHO);
  tcsetattr(STDIN_FILENO, TCSANOW, &newt);
  oldf = fcntl(STDIN_FILENO, F_GETFL, 0);
  fcntl(STDIN_FILENO, F_SETFL, oldf | O_NONBLOCK);

  ch = getchar();

  tcsetattr(STDIN_FILENO, TCSANOW, &oldt);
  fcntl(STDIN_FILENO, F_SETFL, oldf);

  return ch;
}

class ForceServoControlNode : public rclcpp::Node
{
public:
  ForceServoControlNode() : Node("force_servo_control_node")
  {
    // —— 基本参数 ——
    cmd_topic_ = declare_parameter<std::string>("cmd_topic", "force_servo_control_node/delta_twist_cmds");
    ee_link_   = declare_parameter<std::string>("ee_link", "end");  // 末端 link 名字

    // —— PD 控制参数 ——
    kp_ = declare_parameter<double>("kp", 0.1);
    kd_ = declare_parameter<double>("kd", 0.01);
    F_ref_ = declare_parameter<double>("F_ref", 0.8);

    // —— 控制改进 ——
    deadband_ = declare_parameter<double>("deadband", 0.05);
    max_vx_   = declare_parameter<double>("max_vx", 0.05);

    // —— Z 方向恒速运动参数 ——
    z_vel_   = declare_parameter<double>("z_vel", 0.05);   // 速度 m/s
    z_range_ = declare_parameter<double>("z_range", 0.2); // 上下幅度 m

    // 参数动态更新回调
    param_cb_handle_ = this->add_on_set_parameters_callback(
      [this](const std::vector<rclcpp::Parameter> &params) {
        rcl_interfaces::msg::SetParametersResult result;
        result.successful = true;
        for (const auto &p : params) {
          if (p.get_name() == "kp") kp_ = p.as_double();
          else if (p.get_name() == "kd") kd_ = p.as_double();
          else if (p.get_name() == "F_ref") F_ref_ = p.as_double();
          else if (p.get_name() == "deadband") deadband_ = p.as_double();
          else if (p.get_name() == "max_vx") max_vx_ = p.as_double();
          else if (p.get_name() == "z_vel") z_vel_ = p.as_double();
          else if (p.get_name() == "z_range") z_range_ = p.as_double();
        }
        return result;
      }
    );

    // 订阅力数据
    force_sub_ = create_subscription<geometry_msgs::msg::WrenchStamped>(
        "/wrench_comp", 10,
        std::bind(&ForceServoControlNode::force_callback, this, std::placeholders::_1));

    // Publisher
    tf_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
    tf_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf_buffer_);
    twist_cmd_pub_ = create_publisher<geometry_msgs::msg::TwistStamped>(cmd_topic_, 10);

    last_time_ = this->now();

    // 键盘监听线程
    keyboard_thread_ = std::thread([this]() { this->keyboardLoop(); });
  }

  ~ForceServoControlNode()
  {
    stop_keyboard_ = true;
    if (keyboard_thread_.joinable())
      keyboard_thread_.join();
  }

  void initialize()
  {
    setupPlanningSceneMonitor();
    setupServo();
  }

private:
  // ROS 通信
  rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr force_sub_;
  rclcpp::Publisher<geometry_msgs::msg::TwistStamped>::SharedPtr twist_cmd_pub_;
  std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf_listener_;
  std::shared_ptr<planning_scene_monitor::PlanningSceneMonitor> planning_scene_monitor_;
  std::unique_ptr<moveit_servo::Servo> servo_;
  OnSetParametersCallbackHandle::SharedPtr param_cb_handle_;

  // 控制参数
  std::string cmd_topic_;
  std::string ee_link_;
  double kp_{0.01}, kd_{0.01}, F_ref_{0.8};
  double deadband_{0.05};
  double max_vx_{0.06};

  // Z 方向恒速参数
  double z_vel_{0.05};
  double z_range_{0.2};
  double z_pos_{0.0};   // 累积位移
  int z_dir_{1};        // 方向 (+1 上, -1 下)

  // 状态
  double e_last_{0.0};
  rclcpp::Time last_time_;
  bool enabled_{false};

  // 键盘线程
  std::thread keyboard_thread_;
  bool stop_keyboard_{false};

  // 力回调
  void force_callback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
  {
    if (!enabled_) {
      return; // 未启用时直接退出
    }

    double Fz = msg->wrench.force.z;
    double F_meas = Fz;

    double e = F_ref_ - F_meas;

    rclcpp::Time now = (msg->header.stamp.sec > 0 || msg->header.stamp.nanosec > 0)
                         ? rclcpp::Time(msg->header.stamp)
                         : this->now();

    double dt = (now - last_time_).seconds();
    if (dt <= 0.0) dt = 1e-3;
    double de = (e - e_last_) / dt;
    e_last_ = e;
    last_time_ = now;

    // —— X方向力控 ——
    double v_x_cmd = 0.0;
    if (std::fabs(e) > deadband_)
    {
      v_x_cmd = kp_ * e + kd_ * de;
    }
    if (v_x_cmd > max_vx_) v_x_cmd = max_vx_;
    if (v_x_cmd < -max_vx_) v_x_cmd = -max_vx_;

    // —— Z方向恒速上下 ——
    z_pos_ += z_dir_ * z_vel_ * dt;
    if (std::fabs(z_pos_) > z_range_) {
      z_dir_ *= -1;  // 超过范围换向
    }
    double v_z_cmd = z_dir_ * z_vel_;

    // —— 发布控制指令 ——
    auto cmd = std::make_unique<geometry_msgs::msg::TwistStamped>();
    cmd->header.stamp = this->now();
    cmd->header.frame_id = "base_link";
    cmd->twist.linear.x = v_x_cmd;
    cmd->twist.linear.z = v_z_cmd;
    twist_cmd_pub_->publish(*cmd);

    RCLCPP_INFO(LOGGER, "Fz=%.3f e=%.3f vx=%.3f vz=%.3f z_pos=%.3f",
                F_meas, e, v_x_cmd, v_z_cmd, z_pos_);
  }

  void setupPlanningSceneMonitor()
  {
    planning_scene_monitor_ = std::make_shared<planning_scene_monitor::PlanningSceneMonitor>(
        shared_from_this(), "robot_description", tf_buffer_, "planning_scene_monitor");
    if (planning_scene_monitor_->getPlanningScene()) {
      planning_scene_monitor_->startStateMonitor("/joint_states");
    }
  }

  void setupServo()
  {
    auto servo_parameters = moveit_servo::ServoParameters::makeServoParameters(shared_from_this());
    if (!servo_parameters) throw std::runtime_error("Failed to load servo parameters");
    this->set_parameter(rclcpp::Parameter("moveit_servo.publish_period", 0.005));
    this->set_parameter(rclcpp::Parameter("moveit_servo.low_latency_mode", true));
    this->set_parameter(rclcpp::Parameter("moveit_servo.check_collisions", false));
    servo_ = std::make_unique<moveit_servo::Servo>(shared_from_this(), servo_parameters, planning_scene_monitor_);
    servo_->start();
  }

  // 键盘监听循环
  void keyboardLoop()
  {
    RCLCPP_INFO(LOGGER, "Press 'q' to toggle enable/disable force control.");
    while (rclcpp::ok() && !stop_keyboard_)
    {
      int c = getch();
      if (c == 'q' || c == 'Q')
      {
        enabled_ = !enabled_;
        RCLCPP_WARN(LOGGER, "Force control %s", enabled_ ? "ENABLED" : "DISABLED");
      }
      std::this_thread::sleep_for(50ms);
    }
  }
};

int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<ForceServoControlNode>();
  node->initialize();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
